#pragma once

#include <cstdlib>
#include <map>
#include <mutex>
#include "../parallel/lock.h"
#include "auto_close.h"
#include "logger.h"


OPEN_JLIB_NS


template <typename T, std::size_t S = sizeof(T), std::size_t A = alignof(T)>
class AlignAllocator {


public:
  typedef T value_type;
  typedef std::size_t size_type;
  typedef std::ptrdiff_t difference_type;

  typedef T * pointer;
  typedef const T * const_pointer;

  typedef T & reference;
  typedef const T & const_reference;


public:
  inline AlignAllocator () noexcept {
    log_v("AlignAllocator()");
  }

  inline AlignAllocator (const_reference r) noexcept {
    log_v("AlignAllocator(const_reference r = %p)", &r);
  }

  template <typename T2>
  inline AlignAllocator (const AlignAllocator<T2, S, A> &other) noexcept {
    log_v("AlignAllocator(const AlignAllocator<T2> &other = %p)", &other);
  }

  inline ~AlignAllocator () noexcept {
    log_v("~AlignAllocator()");
  }

  inline pointer address (reference r) {
    log_v("address(reference r = %p)", &r);
    return &r;
  }

  inline const_pointer address (const_reference r) const {
    log_v("address(const_reference r = %p)", &r);
    return &r;
  }

  inline pointer allocate (size_type n) {
    log_v("allocate(size_type n = %lu)", n);
    auto p = (size_type)malloc(n * S + A);
    if (!p)
      return nullptr;
    auto u = p % A;
    if (!u)
      return (pointer)p;
    auto a = p + A - u;
    this->map.emplace(a, p);
    return (pointer)a;
  }

  inline void deallocate (pointer p, size_type n) {
    log_v("deallocate(pointer p = %p, size_type n = %lu)", p, n);
    auto a = (size_type)p;
    auto i = this->map.find(a);
    if (i != this->map.end()) {
      p = (pointer)i->second;
      this->map.erase(i);
    }
    free(p);
  }

  inline void construct (pointer p, const_reference v) {
    log_v("construct(pointer p = %p, const_reference v = %p)", p, &v);
    new(p) value_type(v);
  }

  inline void destroy (pointer p) {
    log_v("destroy(pointer p = %p)", p);
    p->~value_type();
  }

  inline size_type max_size () const noexcept {
    log_v("max_size()");
    return ((size_type)-1 - A) / S;
  }

  inline bool operator== (const AlignAllocator<T> &other) const {
    log_v("operator==(const AlignAllocator<T> &other = %p)", &other);
    return this == &other;
  }

  inline bool operator!= (const AlignAllocator<T> &other) const  {
    log_v("operator!=(const AlignAllocator<T> &other = %p)", &other);
    return this != &other;
  }


public:
  template <typename T2>
  struct rebind {
    typedef AlignAllocator<T2> other;
  };


private:
  std::map<size_type, size_type> map;
};






template <typename T, std::size_t S = sizeof(T), std::size_t A = alignof(T)>
class ConcurrentAlignAllocator {


public:
  typedef T value_type;
  typedef std::size_t size_type;
  typedef std::ptrdiff_t difference_type;

  typedef T * pointer;
  typedef const T * const_pointer;

  typedef T & reference;
  typedef const T & const_reference;


public:
  inline ConcurrentAlignAllocator () noexcept {
    log_v("ConcurrentAlignAllocator()");
  }

  inline ConcurrentAlignAllocator (const_reference r) noexcept {
    log_v("ConcurrentAlignAllocator(const_reference r = %p)", &r);
  }

  template <typename T2>
  inline ConcurrentAlignAllocator (const ConcurrentAlignAllocator<T2, S, A> &other) noexcept {
    log_v("ConcurrentAlignAllocator(const ConcurrentAlignAllocator<T2> &other = %p)", &other);
  }

  inline ~ConcurrentAlignAllocator () noexcept {
    log_v("~ConcurrentAlignAllocator()");
  }

  inline pointer address (reference r) {
    log_v("address(reference r = %p)", &r);
    return &r;
  }

  inline const_pointer address (const_reference r) const {
    log_v("address(const_reference r = %p)", &r);
    return &r;
  }

  inline pointer allocate (size_type n) {
    log_v("allocate(size_type n = %lu)", n);
    auto p = (size_type)malloc(n * S + A);
    if (!p)
      return nullptr;
    auto u = p % A;
    if (!u)
      return (pointer)p;
    auto a = p + A - u;
    AutoClose<SpinLock> lock {this->lock};
    this->map.emplace(a, p);
    return (pointer)a;
  }

  inline void deallocate (pointer p, size_type n) {
    log_v("deallocate(pointer p = %p, size_type n = %lu)", p, n);
    this->lock.open();
    auto a = (size_type)p;
    auto i = this->map.find(a);
    if (i != this->map.end()) {
      p = (pointer)i->second;
      this->map.erase(i);
    }
    this->lock.close();
    free(p);
  }

  inline void construct (pointer p, const_reference v) {
    log_v("construct(pointer p = %p, const_reference v = %p)", p, &v);
    new(p) value_type(v);
  }

  inline void destroy (pointer p) {
    log_v("destroy(pointer p = %p)", p);
    p->~value_type();
  }

  inline size_type max_size () const noexcept {
    log_v("max_size()");
    return ((size_type)-1 - A) / S;
  }

  inline bool operator== (const ConcurrentAlignAllocator<T> &other) const {
    log_v("operator==(const ConcurrentAlignAllocator<T> &other = %p)", &other);
    return this == &other;
  }

  inline bool operator!= (const ConcurrentAlignAllocator<T> &other) const  {
    log_v("operator!=(const ConcurrentAlignAllocator<T> &other = %p)", &other);
    return this != &other;
  }


public:
  template <typename T2>
  struct rebind {
    typedef ConcurrentAlignAllocator<T2> other;
  };

private:
  std::map<size_type, size_type> map;
  SpinLock lock;

};


#define ALIGNAS(Type) \
public: \
static jlib::ConcurrentAlignAllocator<Type> __allocator; \
void *operator new (size_t size) { return __allocator.allocate(size); } \
void *operator new[] (size_t size) { return __allocator.allocate(size); } \
void *operator new (size_t size, const std::nothrow_t &) noexcept { return __allocator.allocate(size); } \
void *operator new[] (size_t size, const std::nothrow_t &) noexcept { return __allocator.allocate(size); } \
void operator delete (void *ptr) { __allocator.deallocate((Type*)ptr, 1); } \
void operator delete[] (void *ptr) { __allocator.deallocate((Type*)ptr, 1); } \
void operator delete (void *ptr, const std::nothrow_t &) noexcept { __allocator.deallocate((Type*)ptr, 1); } \
void operator delete[] (void *ptr, const std::nothrow_t &) noexcept { __allocator.deallocate((Type*)ptr, 1); }


CLOSE_JLIB_NS